import os
import random

def split_dataset(path):
    # 先判断存储图像路径和标签的文件是否存在，存在的话就删掉，以免重复录入。
    if os.path.isfile('train.txt') and os.path.isfile('test.txt'):
        os.remove('./train.txt')
        os.remove('./test.txt')
    else:
        pass

    # 存储训练图片的路径和标签
    train_split = []
    # 存储测试图片的路径和标签
    test_split = []
    # 图像标签
    target = 0
    # 存储标签对应类别名称的字典
    dataset_details = {}
    for class_num in os.listdir(path):
        count = 0
        dataset_details[target] = class_num
        class_path = os.path.join(path, class_num)
        for img_name in os.listdir(class_path):
            count += 1
            img_path = os.path.join(class_path, img_name)
            # 每遍历十五张图片就把这张图作为测试图片
            if count % 15 == 0:
                test_split.append(img_path + '\t%d' % target + '\n')
            else:
                train_split.append(img_path + '\t%d' % target + '\n')
        target += 1
    # 打乱数据集
    random.shuffle(train_split)
    random.shuffle(test_split)

    # 将存储列表的内容录入文件里，方便后边加载数据
    with open('./train.txt', mode='a') as f1:
        for train_data in train_split:
            f1.write(train_data)
        f1.close()
    with open('./test.txt', mode='a') as f2:
        for test_data in test_split:
            f2.write(test_data)
        f2.close()

    print(f'训练集图片总数：{len(train_split)}')
    print(f'测试集图片总数：{len(test_split)}')
    return dataset_details


dataset_details = split_dataset(r"C:\pycharmproject\project1\宝石分类\train")
print(f'标签与类别对照:{dataset_details}')